/*
 * TotalEquilibriumProblem.cpp
 *
 *  Created on: 4 Sep 2011
 *      Author: Allan
 */

#include "TotalEquilibriumProblem.h"

TotalEquilibriumProblem::TotalEquilibriumProblem()
{}

TotalEquilibriumProblem::TotalEquilibriumProblem(const Multiphase& multiphase, 
                                                 const ReactionSystem& eReactions, 
                                                 const vector<string>& jSpecies, 
                                                 const vector<string>& eSpecies)
{
	Initialiaze(multiphase, eReactions, jSpecies, eSpecies);
}

TotalEquilibriumProblem::~TotalEquilibriumProblem()
{}

void TotalEquilibriumProblem::Initialiaze(const Multiphase& multiphase, const ReactionSystem& eReactions, const vector<string>& jSpecies, const vector<string>& eSpecies)
{
	// Set data members
	this->multiphase = multiphase;
	this->eReactions = eReactions;
	
	// Set the number of primary (Nj) and equilibrium (Ne) species
	Nj = jSpecies.size();
	Ne = eSpecies.size();
	
	// Set the indexes of the primary (jIndexes) and equilibrium (eIndexes) species
	jIndexes = multiphase[jSpecies];
	eIndexes = multiphase[eSpecies];
	
	// Assemble the stoichiometric matrices (uej) and (uee)
	MatrixXd uej = eReactions.AssembleStoichiometricMatrix(jSpecies);
	MatrixXd uee = eReactions.AssembleStoichiometricMatrix(eSpecies);
	
	// Compute (uee_inv), which is the inverse of the stoichiometric matrix (uee)
	FullPivLU<MatrixXd> lu(uee);
	
	uee_inv = lu.inverse();
	
	// Compute the canonical stoichiometric matrix of the equilibrium-controlled reactions
	vej = - uee_inv * uej;
	
	// Allocate all the auxiliary vector/matrix variables with zero entries
	Ke = VectorXd::Zero(Ne);
	Qe = VectorXd::Zero(Ne);
	ne = VectorXd::Zero(Ne);
}

void TotalEquilibriumProblem::SetParameters(const VectorXd& uj, double T, double P, const VectorXd& n)
{
	this->T  = T;
	this->P  = P;
	this->n  = n;
	this->uj = uj;
	
	// Initialise the equilibrium constants of the equilibrium reactions (Ke)
	Ke = eReactions.EquilibriumConstants(T, P);
	
	// Initialise the vector (ne) with the values from the initial guess (n)
	for(unsigned e = 0; e < Ne; ++e) ne[e] = n[eIndexes[e]];
}

void TotalEquilibriumProblem::Function(const VectorXd& nj, VectorXd& F)
{
	// Update the composition of the equilibrium species
	UpdateEquilibriumComposition(nj);
	
	// Compute the residual vector (F)
	F = nj + vej.transpose() * ne - uj;
}

void TotalEquilibriumProblem::Jacobian(const VectorXd& nj, MatrixXd& J)
{
	// Compute the Jacobian matrix (J)
	J = MatrixXd::Identity(Nj, Nj) + vej.transpose() * ne.asDiagonal() * vej * nj.asDiagonal().inverse();
}

void TotalEquilibriumProblem::UpdateEquilibriumComposition(const VectorXd& nj)
{
	// Check if all the entries in (nj) are non-negative
	if(nj.minCoeff() > 0)
	{
		// Update the nj part of the vector (n) with the values of the vector (nj)
		for(unsigned j = 0; j < Nj; ++j) n[jIndexes[j]] = nj[j];
		
		// Update the activity vector (a) with temperature (T), pressure (P) and the mole vector (n)
		a = multiphase.Activities(T, P, n);
	
		// Update the reaction quotients of the equilibrium reactions (Qe)
		Qe = eReactions.ReactionQuotients(a);
		
		// Update the moles of the equilibrium species (ne)
		const VectorXd lnKe_minus_lnQe = Ke.array().log() - Qe.array().log();
		
		ne = ne.array() * (uee_inv * lnKe_minus_lnQe).array().exp();
		
		// Update the ne part of the vector (n)
		for(unsigned e = 0; e < Ne; ++e) n[eIndexes[e]] = ne[e];
	}
}
